import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader



class StructReconstruct(nn.Module):
    def __init__(self, displace_hidden_size, freqs_size):
        super().__init__()
        self.displace_hidden = displace_hidden_size
        self.freqfeature = nn.Sequential(
            nn.Linear(freqs_size, displace_hidden_size),
            nn.ReLU(),
            nn.Linear(displace_hidden_size, displace_hidden_size)
        )
        self.reconstruct1 = nn.ConvTranspose2d(in_channels=displace_hidden_size, 
                                              out_channels=displace_hidden_size // 4,
                                              kernel_size=(3, 3),
                                              stride=(3, 3))
        self.reconstruct2 = nn.ConvTranspose2d(in_channels=displace_hidden_size // 4, 
                                              out_channels=displace_hidden_size // 16,
                                              kernel_size=(3, 3), 
                                              stride=(3, 3))
        self.reconstruct3 = nn.ConvTranspose2d(in_channels=displace_hidden_size // 16, 
                                              out_channels=2,
                                              kernel_size=(5, 1),
                                              stride=(5, 1))
        self.relu = nn.ReLU()

    def forward(self, displace_hidden, freqs):
        feature = displace_hidden + self.freqfeature(freqs)
        # feature = self.freqfeature(freqs)
        feature = feature.unsqueeze(-1).unsqueeze(-1)
        output = self.reconstruct1(feature)
        # output = self.relu(output)
        output = self.reconstruct2(output)
        # output = self.relu(output)
        output = self.reconstruct3(output)

        return output


